package problems;

import java.util.Random;

import globals.CentralStatics;

public class Momdp<S,A> {

	public static final Integer TERMINAL_STATE = new Integer(-1);
	public static final Integer UNSPECIFIED_STATE = new Integer(-2);
	
	protected boolean finiteHorizon;
	protected int horizon;
	protected double discountFactor;
	
	protected int currentState;
	
	protected EnumerableStates<S> states;
	protected TransitionFunction<S,A> transitions; 
	protected Actions<A> actions;
	protected RewardFunction<S,A> rewards;
	private String problemName;
	
	protected double[] initialStateDistribution;
	
	public Momdp(EnumerableStates<S> s_, TransitionFunction<S,A> t_, Actions<A> a_,
			RewardFunction<S,A> r_, double[] mu, double g_){
		states = s_;
		transitions = t_;
		actions = a_;
		rewards = r_;
		
		if(mu != null){
			this.reInitialize();
		} else {
			mu = new double[this.states.numberOfStates()];
		}
		
		this.initialStateDistribution = mu;
		
		this.problemName = ""+CentralStatics.getNextProblemIdentifier()+": ";
		
		this.finiteHorizon = false;
		this.horizon = Integer.MAX_VALUE;
		this.discountFactor = g_;
	}
	
	public void reInitialize(){
		Random rand = CentralStatics.getCentralRandom();
		double stateDbl = rand.nextDouble();
		double curr = 0;
		this.currentState = -1;
		for(int i=0; i<this.initialStateDistribution.length; i++){
			curr+=initialStateDistribution[i];
			if(stateDbl<=curr){
				this.currentState = i;
				break;
			}
			if(i==(initialStateDistribution.length-1) && this.currentState==-1){
				//Fix for precision errors of doubles; 
				//beliefs might not exactly sum to 1;
				this.currentState = i;
			}
		}
		if(currentState==-1){
			System.err.println("(re)initialization of MOMDP failed");
			System.exit(-1);
		}
	}
	
	public S currentState(){
		return this.states.stateIdentifier(this.currentState);
	}
	
	public int currentStateNumber(){
		return (this.currentState);
	}
	
	public double[] performAction(A action){
		if(currentState==TERMINAL_STATE){
			return null;
		}
		if(currentState==UNSPECIFIED_STATE){
			System.err.println("Trying to simulate action in unspecified state.");
			System.exit(UNSPECIFIED_STATE);
			return null;
		}
		S currentS  = states.stateIdentifier(currentState);
		S nextState = this.transitions.nextState(currentS, action);
		this.currentState = states.stateNumber(nextState);
		double[] reward = this.rewards.getReward(currentS, action, nextState);
		if(this.inTerminalState()){
			this.reInitialize();
			return reward;
		} else {
			return reward;
		}
	}
	
	public double[] reward(S state, A action){
		return this.rewards.expectedReward(state, action);
	}
	
	public double[] reward(int state, A action){
		S s = this.states.stateIdentifier(state);
		return this.rewards.expectedReward(s, action);
	}
	
	public Actions<A> getActionSet(){
		return this.actions;
	}
	
	public EnumerableStates<S> getStateSet(){
		return this.states;
	}
	
	public A randomAction(){
		return this.actions.randomAction();
	}
	
	public String toString(){
		String result ="MOMDP "+this.problemName+"\n";
		int noS = this.states.numberOfStates();
		result+= ""+noS+" states; currently in state "+this.currentState+"\n";
		return result;
	}
	
	public int numberOfObjectives(){
		return this.rewards.numberOfObjectives();
	}
	
	public int numberOfStates(){
		return this.states.numberOfStates();
	}
	
	public boolean inTerminalState(){
		return this.states.isTerminal(this.currentState);
	}
	
	public double transitionProb(S s, A a, S s_){
		return this.transitions.probabilityOfTransition(s, a, s_);
	}
	
	public double getGamma(){
		return this.discountFactor;
	}
	
	public boolean isDiscounted(){
		boolean result = this.discountFactor<1;
		return result;
	}
	
}
